# encoding:utf-8
import cv2
import numpy as np
import requests
import base64
import glob
import os


# args: 
# seg_dir: dir of segmentation images(4 channels)
# suggest_size: the suggested size, final size will not exceed this setting
def AdjustCropAndResize(seg_dir, suggest_size):
    begin_idx = 0
    out_path = './CROP/'
    if os.path.exists(out_path):
        files_crop = glob.glob(out_path + '/*.PNG')
        for i in range(len(files_crop)):
            id_ = int(files_crop[i][-12:-4]) + 1
            begin_idx = max(id_, begin_idx)
    else:
        os.mkdir(out_path)
    print(begin_idx)
    
    files_seg = glob.glob(seg_dir + '/*')
    for i in range(len(files_seg)):
        print(files_seg[i])
        im_ = cv2.imread(files_seg[i], -1)
        assert len(im_.shape)==3 and im_.shape[2]==4
        im_fg = im_[:,:,:3]
        im_seg = im_[:,:,3]
        w, h = im_seg.shape[1], im_seg.shape[0]
        cast_x = np.sum(im_seg, axis=0, keepdims=False)>0
        cast_y = np.sum(im_seg, axis=1, keepdims=False)>0
        x1 = np.argmax(cast_x)
        x2 = w - np.argmax(cast_x[::-1])
        y1 = np.argmax(cast_y)
        y2 = h - np.argmax(cast_y[::-1])
        # add some protecting pixels around target region
        protect_size = 0
        x1 = max(0, x1-protect_size)
        x2 = min(w, x2+protect_size)
        y1 = max(0, y1-protect_size)
        y2 = min(h, y2+protect_size)
        # crop this region
        im_seg = im_seg[y1:y2, x1:x2]
        im_fg = im_fg[y1:y2, x1:x2]
        h, w = im_seg.shape[0:2]
        max_w = suggest_size[0]
        max_h = suggest_size[1]
        ratio_x = 1.0*max_w / w
        ratio_y = 1.0*max_h / h
        if ratio_x < 1 or ratio_y < 1:
            ratio_xy = min(ratio_x, ratio_y)
        w_new = int(w * ratio_xy)
        h_new = int(h * ratio_xy)
        im_seg = cv2.resize(im_seg, (w_new, h_new))
        im_fg = cv2.resize(im_fg, (w_new, h_new))
        im_rgbd = np.stack([im_fg[:,:,0],
                            im_fg[:,:,1],
                            im_fg[:,:,2],
                            im_seg], axis=-1)
        cv2.imwrite(out_path + '/%08d.PNG'%(i+begin_idx), im_rgbd)



if __name__ == '__main__':
    seg_dir = './OUTPUT-ALL'
    AdjustCropAndResize(seg_dir, [512, 512])
    